import numpy as np
import matplotlib.pyplot as plt


class PIDController:
    def __init__(self, Kp, Ki, Kd, integrator_max=None, integrator_min=None,max=None,min=None):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.integrator = 0
        self.prev_error = 0
        self.integrator_max = integrator_max
        self.integrator_min = integrator_min
        self.max = max
        self.min = min
        self.error = 0
        self.measurement_history = []
        self.setpoint_history = []

    def reset(self):
        self.integrator = 0
        self.prev_error = 0

    def update(self, setpoint, measurement, dt):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        error = setpoint - measurement
        self.error = error
        self.integrator += error * dt
        if self.integrator_max is not None and self.integrator > self.integrator_max:
            self.integrator = self.integrator_max
        if self.integrator_min is not None and self.integrator < self.integrator_min:
            self.integrator = self.integrator_min

        derivative = (error - self.prev_error) / dt
        output = self.Kp * error + self.Ki * self.integrator + self.Kd * derivative
        self.prev_error = error
        if self.max is not None and self.min is not None:
            output = np.clip(output,self.min,self.max)
        return output

    def plot_error(self,fig,i,title='none'):
        # plt.plot(self.setpoint_history, [i for i in self.setpoint_history], linewidth =2.0, label = r"setpoint",color='b', linestyle='dotted')
        # plt.plot(self.measurement_history, [i for i in self.measurement_history], linewidth =2.0, label = r"setpoint",color='r',linestyle='dashed')
        ax = fig.add_subplot(131 + i)
        ax.plot([i/100 for i in range(len(self.setpoint_history))], self.setpoint_history, label='setpoint')
        ax.plot([i/100 for i in range(len(self.measurement_history))], self.measurement_history,label='measurement')
        ax.set_xlabel('time')
        ax.set_ylabel('value')
        ax.set_title(title)
        ax.legend()


class CascadePIDController:
    def __init__(self, outer_controller, inner_controller):
        self.outer_controller = outer_controller
        self.inner_controller = inner_controller

    def update(self, outer_setpoint, outer_measurement, inner_measurement, dt):
        outer_output = self.outer_controller.update(outer_setpoint, outer_measurement, dt)
        inner_output = self.inner_controller.update(outer_output, inner_measurement, dt)
        return inner_output


class posion_pos_controller:
    def __init__(self,params):
        self.params = params

    def position_controller(self):
        x_position_controller = PIDController(self.params.Kpxp, 0, 0, integrator_max=1,
                                              integrator_min=-1, max=5, min=-5)
        x_velocity_controller = PIDController(self.params.Kvxp, self.params.Kvxi, self.params.Kvxd, integrator_max=1,
                                              integrator_min=-1)
        y_position_controller = PIDController(self.params.Kpyp, 0, 0, integrator_max=1,
                                              integrator_min=-1, max=5, min=-5, )
        y_velocity_controller = PIDController(self.params.Kvyp, self.params.Kvyi, self.params.Kvyd, integrator_max=1,
                                              integrator_min=-1)

        height_position_controller = PIDController(self.params.Kpzp, 0, 0, integrator_max=1,
                                                   integrator_min=-1, min=-3, max=3)
        height_velocity_controller = PIDController(self.params.Kvzp, self.params.Kvzi, self.params.Kvzd, integrator_max=1,
                                                   integrator_min=-1, min=-0.4, max=0.4)
        x_controller = CascadePIDController(x_position_controller, x_velocity_controller)
        y_controller = CascadePIDController(y_position_controller, y_velocity_controller)
        height_controller = CascadePIDController(height_position_controller, height_velocity_controller)
        return [x_controller,y_controller,height_controller]

    def pos_controller(self):
        DEG2RAD = self.params.DEG2RAD
        roll_angle_controller = PIDController(self.params.Kp_RP_ANGLE, 0, 0, integrator_max=1,
                                              integrator_min=-1, max=DEG2RAD * 35, min=-DEG2RAD * 35)
        roll_rate_controller = PIDController(self.params.Kp_RP_AngleRate, self.params.Ki_RP_AngleRate,
                                             self.params.Kd_RP_AngleRate, integrator_max=1,
                                             integrator_min=-1, max=DEG2RAD * 220, min=-DEG2RAD * 220)
        roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)

        pitch_angle_controller = PIDController(self.params.Kp_RP_ANGLE, 0, 0, integrator_max=1,
                                               integrator_min=-1, max=DEG2RAD * 35, min=-DEG2RAD * 35)
        pitch_rate_controller = PIDController(self.params.Kp_RP_AngleRate, self.params.Ki_RP_AngleRate,
                                              self.params.Kd_RP_AngleRate, integrator_max=1,
                                              integrator_min=-1, max=DEG2RAD * 220, min=-DEG2RAD * 220)
        pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)
        yaw_angle_controller = PIDController(self.params.Kp_RP_ANGLE, 0, 0, integrator_max=1,
                                             integrator_min=-1, max=DEG2RAD * 35, min=-DEG2RAD * 35)
        yaw_rate_controller = PIDController(self.params.Kp_YAW_AngleRate, self.params.Ki_YAW_AngleRate,
                                            self.params.Kd_YAW_AngleRate, integrator_max=1,
                                            integrator_min=-1, max=DEG2RAD * 220, min=-DEG2RAD * 220)
        yaw_controller = CascadePIDController(yaw_angle_controller, yaw_rate_controller)
        return [roll_controller,pitch_controller,yaw_controller]
